#pragma once

#include "boost/context/fcontext.hpp"
#include <memory>
#include <deque>
#include <string>
#include <functional>
#include <iostream>
#include <unordered_set>
#include <thread>
#include <chrono>
#include <algorithm>

struct CODispatcher;

struct COTask
{
	enum
	{
		STATUS_NONE,
		STATUS_WAIT,
		STATUS_RUNNING,
		STATUS_FINISH,
	};

	COTask(CODispatcher* dispatcher, std::size_t stack_size, const std::function<void()>& func);
	~COTask();

	static void entry(intptr_t v);

	inline void run()
	{
		this->m_func();
	}

	void handle_finish();
private:
	friend struct CODispatcher;
private:
	char* m_stack;
	unsigned char m_status;
	std::size_t m_stack_size;
	boost::context::fcontext_t m_context;
	CODispatcher* m_dispatcher;
	std::function<void()> m_func;
};

struct CODispatcher
{
	CODispatcher();
	~CODispatcher();
    COTask* create(const std::function<void()>& func, std::size_t stack_size = 50* 1024);

	void yield();
	void suspend();
	void resume(COTask* t);
	void loop_once();
	void loop(unsigned int sleep_microsecs = 100);
	void cleanup();

	inline void operator -(const std::function<void()>& func)
	{
		this->create(func);
	}
	inline bool intask()
	{
		return m_current != nullptr;
	}
	inline COTask* current()
	{
		return m_current;
	}
	inline bool empty()
	{
		return m_tasks.empty();
	}

	static void startup_current_thread();
	static void cleanup_current_thread(bool release_dispatcher = true);
private:
	COTask* m_current;
	boost::context::fcontext_t m_main_context;
	std::deque<COTask*> m_ready_tasks;
	std::deque<COTask*> m_finished_tasks;
	std::unordered_set<COTask*> m_tasks;
};

extern CODispatcher* get_current_codispatcher();

struct COSelector;

template<typename T> struct COChan
{
	COChan<T>& operator <<(const T& v)
	{
		this->push(v);
		return *this;
	}

	COChan<T>& operator >>(T& v)
	{
		this->pop(v);
		return *this;
	}

	bool try_pop(T& v)
	{
		if (!m_queue.empty())
		{
			v = std::move(m_queue.front());
			m_queue.pop_front();
			return true;
		}
		return false;
	}

	bool pop(T& v)
	{
		auto dispatcher = get_current_codispatcher();

		while (true)
		{
			if (!m_queue.empty())
			{
				v = std::move(m_queue.front());
				m_queue.pop_front();
				return true;
			}
			if (dispatcher->intask())
			{
				m_rtasks.push_back(dispatcher->current());
				dispatcher->suspend();
			}
			else
			{ 
				break;
			}
		}
		return false;
	}

	void push(const T& v)
	{
		m_queue.push_back(v);

		while (!m_rtasks.empty())
		{
			auto dispatcher = get_current_codispatcher();
			dispatcher->resume(m_rtasks.front());
			m_rtasks.pop_front();
		}

		for (auto& v : m_rwaiters)
		{
			(*v.chan) << v.index;
		}
	}
	
	typedef std::function<void(COSelector& receive, std::size_t, COChan<std::size_t>&)> ReceiverRegisterFunc;

	std::tuple<std::function<void()>, ReceiverRegisterFunc, std::function<void(COSelector&)>> operator>>(const std::function<void()>& func)
	{
		return std::make_tuple(func, [this](COSelector& selector, std::size_t index, COChan<std::size_t>& wait_chan){
			SWaiter waiter;
			waiter.selector = &selector;
			waiter.index = index;
			waiter.chan = &wait_chan;
			m_rwaiters.push_back(waiter);
		}, [this](COSelector& selector) {
			this->cancel_receiver(selector);
		});
	}

	void cancel_receiver(COSelector& selector)
	{
		m_rwaiters.erase(std::remove_if(m_rwaiters.begin(), m_rwaiters.end(), [&](SWaiter& r){ return r.selector == &selector; }), m_rwaiters.end());
	}

	struct SWaiter
	{
		COSelector* selector;
		COChan<std::size_t>* chan;
		std::size_t index;
	};

	std::deque<T> m_queue;
	std::deque<COTask*> m_rtasks;
	std::deque<SWaiter> m_rwaiters;
};

struct COSelector
{
	bool m_canceled;
	COChan<std::size_t> wait_chan;
	std::vector<std::function<void()>> wait_funcs;
	std::vector<std::function<void(COSelector&)>> cancel_funcs;

	COSelector()
		: m_canceled(false)
	{		
	}

	void exit()
	{
		if (m_canceled)
			return;
		m_canceled = true;
		for (auto& func : cancel_funcs)
			func(*this);
		wait_chan.push(std::size_t(-1));
	}

	template<typename... T> void loop(T&&... args)
	{
		co_select(*this, args...);
	}
};

extern void _co_select(COSelector& selector);

template<typename HEAD, typename... T> void _co_select(COSelector& selector, HEAD&& head, T&&... args)
{
	std::function<void()> func;
	std::function<void(COSelector&, std::size_t,COChan<std::size_t>&)> register_func;
	std::function<void(COSelector&)> cancel_func;
	std::tie(func, register_func, cancel_func) = head;

	auto index = selector.wait_funcs.size();
	selector.wait_funcs.push_back(func);
	register_func(selector, index, selector.wait_chan);
	selector.cancel_funcs.push_back(cancel_func);
	_co_select(selector, args...);
}

template<typename... T> void co_select(COSelector& selector, T&&... args)
{
	_co_select(selector, args...);
	while (true)
	{
		size_t i = 0;
		selector.wait_chan >> i;
		if (i == std::size_t(-1))
			return;
		selector.wait_funcs[i]();
	}
}

template<typename... T> void co_select(T&&... args)
{
	COSelector selector;
	co_select(selector, args...);
}

#define co_yield get_current_codispatcher()->yield()
#define co_suspend get_current_codispatcher()->suspend()
#define co_resume(t) get_current_codispatcher()->resume(t)
#define co_current() get_current_codispatcher()->current()
#define co_sleep(micro_seconds)
#define co_loop_main while(true) { get_current_codispatcher()->loop_once(); std:::this_thread::sleep_for(std::chrono::microseconds * 100); }
#define co_loop_once get_current_codispatcher()->loop_once()
#define co_loop get_current_codispatcher()->loop()
#define co_wait_events while(true) co_suspend
#define go (*get_current_codispatcher()) - 
